# 指定文件编码为UTF-8
# coding: utf-8

"""
optimizer_compare_mnist.py 功能说明：
1. 比较不同优化算法在MNIST数据集上的表现
2. 测试的优化器包括：SGD、Momentum、AdaGrad、Adam
3. 使用相同的4层神经网络结构进行公平比较
4. 记录并可视化训练过程中的损失变化
5. 使用平滑曲线展示训练趋势

关键点：
- 直观展示不同优化器的收敛特性
- 使用相同的网络结构和超参数设置
- 批量大小固定为128
- 最大迭代次数2000次
- 使用平滑曲线处理噪声数据
"""

# 导入操作系统模块
import os
# 导入系统模块
import sys
# 添加父目录到系统路径，以便导入父目录中的模块
sys.path.append(os.pardir)
# 导入matplotlib绘图库
import matplotlib.pyplot as plt
# 从dataset.mnist导入MNIST数据加载函数
from dataset.mnist import load_mnist
# 从common.util导入平滑曲线函数
from common.util import smooth_curve
# 从common.multi_layer_net导入多层网络类
from common.multi_layer_net import MultiLayerNet
# 从common.optimizer导入所有优化器
from common.optimizer import *

# 0.加载MNIST数据 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
# 加载MNIST数据集(自动归一化)
(x_train, t_train), (x_test, t_test) = load_mnist(normalize=True)

# 设置训练参数
train_size = x_train.shape[0]  # 训练数据大小
batch_size = 128               # 批量大小
max_iterations = 2000          # 最大迭代次数

# 1.实验设置 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
# 初始化各种优化器
optimizers = {}
optimizers['SGD'] = SGD()           # 随机梯度下降
optimizers['Momentum'] = Momentum() # 动量法
optimizers['AdaGrad'] = AdaGrad()   # AdaGrad
optimizers['Adam'] = Adam()         # Adam
#optimizers['RMSprop'] = RMSprop()  # RMSprop(注释状态)

# 初始化网络和损失记录
networks = {}      # 存储各优化器对应的网络
train_loss = {}    # 存储各优化器的损失记录

# 为每个优化器创建相同的网络结构
for key in optimizers.keys():
    # 4层神经网络(784-100-100-100-100-10)
    networks[key] = MultiLayerNet(
        input_size=784,
        hidden_size_list=[100, 100, 100, 100],
        output_size=10)
    # 初始化损失记录列表
    train_loss[key] = []

# 2.训练过程 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
for i in range(max_iterations):
    # 随机选择批量数据
    batch_mask = np.random.choice(train_size, batch_size)
    x_batch = x_train[batch_mask]  # 批量输入数据
    t_batch = t_train[batch_mask]  # 批量标签数据

    # 对每个优化器进行参数更新
    for key in optimizers.keys():
        # 计算梯度
        grads = networks[key].gradient(x_batch, t_batch)
        # 更新参数
        optimizers[key].update(networks[key].params, grads)
        # 计算并记录损失
        loss = networks[key].loss(x_batch, t_batch)
        train_loss[key].append(loss)

    # 每100次迭代打印进度
    if i % 100 == 0:
        print("===========" + "iteration:" + str(i) + "===========")
        # 打印各优化器的当前损失
        for key in optimizers.keys():
            loss = networks[key].loss(x_batch, t_batch)
            print(key + ":" + str(loss))

# 3.结果可视化 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
# 设置各优化器的绘图标记样式
markers = {"SGD": "o", "Momentum": "x", "AdaGrad": "s", "Adam": "D"}
# 创建x轴数据(迭代次数)
x = np.arange(max_iterations)

# 绘制各优化器的损失曲线
for key in optimizers.keys():
    # 使用平滑曲线绘制，每100个点标记一次
    plt.plot(x, smooth_curve(train_loss[key]),
             marker=markers[key],
             markevery=100,
             label=key)

# 设置图表属性
plt.xlabel("iterations")  # x轴标签
plt.ylabel("loss")        # y轴标签
plt.ylim(0, 1)           # y轴范围
plt.legend()             # 显示图例
plt.show()               # 显示图表

"""
预期结果分析：
1. SGD: 收敛较慢，可能震荡明显
2. Momentum: 收敛较快，震荡较小
3. AdaGrad: 初期收敛快，后期可能变慢
4. Adam: 通常表现最好，收敛快且稳定

注意事项：
1. 每次运行结果可能略有不同
2. 可以尝试调整学习率等超参数
3. 增加迭代次数可能看到更明显的差异
4. 不同问题的最佳优化器可能不同
"""
